from comman.functions import *


class SoftmaxWithLoss:
    def __init__(self):
        self.loss = None
        self.y = None  # softmax的输出
        self.t = None  # 监督数据

    def forward(self, x, t):
        self.t = t
        self.y = softmax(x)
        self.loss = cross_entropy_error(self.y, self.t)

        return self.loss

    def backward(self, dout=1):
        batch_size = self.t.shape[0]
        if self.t.size == self.y.size:  # 监督数据是one-hot-vector的情况
            dx = (self.y - self.t) / batch_size
        else:
            dx = self.y.copy()
            dx[np.arange(batch_size), self.t] -= 1
            dx = dx / batch_size

        return dx * dout


class MeanSquaredLoss:
    def __init__(self):
        self.x = None  # 网络输出
        self.t = None  # 监督数据

    def forward(self, x, t):
        self.x = x
        self.t = t

        y = np.square(x - t)
        loss = np.mean(y)

        return loss

    def backward(self, dout=1):
        batch_size = self.t.shape[0]
        if self.t.size == self.x.size:  # 监督数据是one-hot-vector的情况
            dx = self.x - self.t
        else:
            dx = self.x.copy()
            dx[np.arange(batch_size), self.t] -= 1
        dx = dx / batch_size
        return dx * dout

